import os
from flask import Flask, request, jsonify, send_file
from selenium import webdriver
from selenium.webdriver.firefox.options import Options
from selenium.webdriver.common.by import By
import time

app = Flask(__name__)

# 确保文件保存目录存在
OUTPUT_DIR = "D:/file"
if not os.path.exists(OUTPUT_DIR):
    os.makedirs(OUTPUT_DIR)


def convert_url_to_pdf(url, file_path):
    # 设置 Firefox 为 headless 模式
    options = Options()
    options.headless = True

    # 设置 Firefox 打印配置（以 PDF 保存）
    prefs = {
        "printing.print_preview_sticky_settings.appstate": "false",
        "savefile.default_directory": OUTPUT_DIR,
        "print.always_print_silent": True,
        "print.show_print_progress": False
    }
    options.set_preference("browser.printing.print_preview", False)
    options.set_preference("browser.printing.print_printer", "PDF")
    options.set_preference("browser.download.folderList", 2)  # 使用自定义文件夹
    options.set_preference("browser.download.dir", OUTPUT_DIR)  # 设置文件保存路径

    # 启动 Firefox 浏览器
    driver = webdriver.Firefox(options=options)

    try:
        # 导航到指定 URL
        driver.get(url)

        # 触发页面打印
        driver.execute_script('window.print();')

        # 稍等一下让文件保存完成
        time.sleep(5)  # 等待文件生成

    finally:
        driver.quit()


@app.route('/convert', methods=['GET'])
def convert():
    # 从请求中获取 URL 参数和文件名参数
    url = request.args.get('url')
    file_name = request.args.get('file_name', 'output.pdf')  # 默认文件名为 output.pdf

    if not url:
        return jsonify({"error": "URL parameter is required"}), 400

    # 设置文件保存路径
    file_path = os.path.join(OUTPUT_DIR, file_name)

    try:
        # 转换网页为 PDF 并保存
        convert_url_to_pdf(url, file_path)

        # 返回文件的下载链接
        return jsonify({"message": f"PDF saved successfully", "file_path": file_path})

    except Exception as e:
        return jsonify({"error": str(e)}), 500


@app.route('/download/<filename>', methods=['GET'])
def download(filename):
    # 确保文件存在
    file_path = os.path.join(OUTPUT_DIR, filename)

    if os.path.exists(file_path):
        return send_file(file_path, as_attachment=True)
    else:
        return jsonify({"error": "File not found"}), 404


if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)
